import cv2
import numpy as np
from ultralytics import YOLO


def load_yolov8_model():
    # 加载YOLOv8模型 (会自动下载预训练权重如果本地没有)
    model = YOLO('yolov8n.pt')  # 可以使用'yolov8s.pt', 'yolov8m.pt'等不同大小的模型
    return model


def draw_detections(frame, results, class_names, colors):
    # 遍历所有检测结果
    for result in results:
        # 获取边界框坐标、类别ID和置信度
        boxes = result.boxes.xyxy.cpu().numpy()
        class_ids = result.boxes.cls.cpu().numpy().astype(int)
        confidences = result.boxes.conf.cpu().numpy()

        # 遍历每个检测结果
        for box, class_id, confidence in zip(boxes, class_ids, confidences):
            # 只处理人的类别 (COCO数据集中人的类别ID是0)
            if class_id == 0:
                x1, y1, x2, y2 = map(int, box)

                # 获取类别标签和颜色
                label = class_names[class_id]
                color = colors[class_id]

                # 绘制边界框
                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

                # 准备标签文本 (类别+置信度)
                text = f"{label}: {confidence:.2f}"

                # 计算标签文本的大小
                (text_width, text_height), baseline = cv2.getTextSize(
                    text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)

                # 绘制标签背景矩形
                cv2.rectangle(frame, (x1, y1 - text_height - baseline),
                              (x1 + text_width, y1), color, -1)

                # 绘制标签文本
                cv2.putText(frame, text, (x1, y1 - baseline),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

    return frame


def main():
    # 加载COCO数据集类别名 (80个类别)
    with open("coco.names", "r") as f:
        class_names = [line.strip() for line in f.readlines()]

    # 为不同类别生成随机颜色
    colors = np.random.uniform(0, 255, size=(len(class_names), 3))

    # 加载YOLOv8模型
    model = load_yolov8_model()

    # 打开摄像头
    cap = cv2.VideoCapture(0)

    while True:
        # 读取一帧图像
        ret, frame = cap.read()
        if not ret:
            break

        # 使用YOLOv8进行检测
        results = model(frame, verbose=False)  # verbose=False关闭控制台输出

        # 绘制检测结果
        frame = draw_detections(frame, results, class_names, colors)

        # 显示结果
        cv2.imshow("YOLOv8 Object Detection", frame)

        # 按'q'键退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放资源
    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()